// ComfyUI.mxToolkit.Reroute v.0.9e - Max Smirnov 2024
import { app } from "../../scripts/app.js";
import { mergeIfValid, getWidgetConfig, setWidgetConfig } from "../core/widgetInputs.js";

app.registerExtension({
	name: "Comfy.MxRerouteNode",
	registerCustomNodes(app) {
		class MxRerouteNode {
			constructor() {
				if (!this.properties) {
					this.properties = {};
				}
				this.addInput("", "*");
				this.addOutput("", "*");

				this.properties.inputDir = "LEFT";
				this.properties.outputDir= "RIGHT";
				this.linkType = "*";
				this.bgcolor="rgba(20,20,200,0.0)";
				this.keyCode = 0;

				this.size = [2*LiteGraph.CANVAS_GRID_SIZE + 1.4*LiteGraph.NODE_SLOT_HEIGHT, 2*LiteGraph.CANVAS_GRID_SIZE + 1.4*LiteGraph.NODE_SLOT_HEIGHT];

				this.onAdded = function ()
				{
					this.resizable = false;
					this.ioOrientation();
				}

				this.onDrawBackground = function (ctx)
				{
					if (!this.configured) { this.configured = true; this.onConnectionsChange(); }
					const canvas = app.graph.list_of_graphcanvas[0];
					let linkColor = LGraphCanvas.link_type_colors[this.linkType];
					if (linkColor === "") linkColor = LiteGraph.LINK_COLOR;

					if (this.inputs && this.outputs) if (this.inputs[0].pos && this.outputs[0].pos)
					{
						const drawLink = (ctx) =>
						{
							ctx.beginPath();
							ctx.moveTo(this.inputs[0].pos[0], this.inputs[0].pos[1]);
							if (canvas.links_render_mode < 2)
							{
								ctx.lineTo(this.size[0]/2, this.size[1]/2);
								ctx.lineTo(this.outputs[0].pos[0], this.outputs[0].pos[1]);
							} else ctx.quadraticCurveTo(this.size[0]/2, this.size[1]/2, this.outputs[0].pos[0], this.outputs[0].pos[1]);
							ctx.stroke();
						}
						this.inputs[0].color_on = "rgba(0,0,0,0)";
						this.outputs[0].color_on = "rgba(0,0,0,0)";
						this.outputs[0].shape = LiteGraph.ROUND_SHAPE;
						if (this.outputs[0].links) if (this.outputs[0].links.length) if (this.outputs[0].links.length > 0) this.outputs[0].shape = LiteGraph.GRID_SHAPE;

					    if (canvas) if (canvas.render_connections_border && canvas.ds.scale > 0.6)
					    {
							ctx.lineWidth = canvas.connections_width + 4;
							ctx.strokeStyle = "rgba(0,0,0,0.5)";
							drawLink(ctx);
					    }
 						ctx.lineWidth = canvas.connections_width;
 						ctx.strokeStyle = linkColor;
 						drawLink(ctx);
					}

					if (this.mouseOver) if (this.inputs && this.outputs) if (this.inputs[0].pos && this.outputs[0].pos)
					{
						ctx.lineWidth = 1;
 						ctx.strokeStyle = linkColor;
						ctx.beginPath();
						ctx.moveTo(this.pos[0], this.pos[1]);
						ctx.roundRect(0, 0, this.size[0], this.size[1], 10);
						ctx.stroke();
						ctx.fillStyle = linkColor;
						ctx.beginPath();
            			ctx.arc(this.inputs[0].pos[0], this.inputs[0].pos[1], 5, 0, 2 * Math.PI, false);
            			ctx.arc(this.outputs[0].pos[0], this.outputs[0].pos[1], 5, 0, 2 * Math.PI, false);
            			ctx.fill();
					}
				}

				this.onKeyUp = function(e)
				{
					if (e.keyCode < 37 || e.keyCode > 40) return;
					if (this.keyCode > 0)
					{
    					const arrowKeys = { 37: "LEFT", 38: "UP", 39: "RIGHT", 40: "DOWN"};
    					this.properties.inputDir  = arrowKeys[this.keyCode];
    					this.properties.outputDir = arrowKeys[e.keyCode];
    					this.onPropertyChanged();
    					this.keyCode = 0;
					} else this.keyCode = e.keyCode;
				}

				this.onDeselected = () => {this.keyCode = 0};

				this.getExtraMenuOptions = function()
				{
				    const that = this;
				    const iDir = ["LEFT",  "RIGHT", "DOWN", "UP",   "LEFT", "LEFT", "RIGHT", "RIGHT", "UP",    "DOWN",  "DOWN", "UP"];
				    const oDir = ["RIGHT", "LEFT",  "UP",   "DOWN", "UP",   "DOWN", "UP",    "DOWN",  "RIGHT", "RIGHT", "LEFT", "LEFT"];
				    const key  = ["🠖",     "🠔",    "🠕",    "🠗",    "⮥",    "⮧",    "⮤",    "⮦",     "⮡",     "⮣",     "⮢",   "⮠"], options = [];
				    for (let i = 0; i < 12; i++) options.push({ content: key[i], callback: () => { that.properties.inputDir = iDir[i]; that.properties.outputDir = oDir[i]; that.onPropertyChanged(); }});
				    return options;
				};

				this.onPropertyChanged = function ()
				{
					const aValues = ["LEFT","RIGHT","UP","DOWN","TOP"];
					const sValues = ["L","R","U","D","T"]
					this.properties.inputDir = this.properties.inputDir.toUpperCase();
					this.properties.outputDir = this.properties.outputDir.toUpperCase();
					if (sValues.indexOf(this.properties.inputDir)>=0) this.properties.inputDir = aValues[sValues.indexOf(this.properties.inputDir)];
					if (sValues.indexOf(this.properties.outputDir)>=0) this.properties.outputDir = aValues[sValues.indexOf(this.properties.outputDir)];
					if (this.properties.inputDir === "TOP") this.properties.inputDir = "UP";
					if (this.properties.outputDir === "TOP") this.properties.inputDir = "UP";
					this.properties.inputDir = aValues.includes(this.properties.inputDir)?this.properties.inputDir:"LEFT";
					this.properties.outputDir= aValues.includes(this.properties.outputDir)?this.properties.outputDir:"RIGHT";
					if (this.properties.inputDir === this.properties.outputDir) this.properties.outputDir=(this.properties.outputDir==="RIGHT")?"LEFT":"RIGHT";
					this.ioOrientation();
					app.graph.setDirtyCanvas(true, true);
				}

				this.ioOrientation = function ()
				{
					let i = this.properties.inputDir;
					let o = this.properties.outputDir;
            	    switch (i)
            	    {
            	        case "LEFT":
            	            this.inputs[0].pos = [0,0.7*LiteGraph.NODE_SLOT_HEIGHT+LiteGraph.CANVAS_GRID_SIZE];
            	            this.inputs[0].dir = LiteGraph.LEFT;
            	            break;
            	        case "RIGHT":
            	        	this.inputs[0].pos = [this.size[0],0.7*LiteGraph.NODE_SLOT_HEIGHT+LiteGraph.CANVAS_GRID_SIZE];
            	            this.inputs[0].dir = LiteGraph.RIGHT;
            	            break;
            	        case "UP":
            	        	this.inputs[0].pos = [this.size[0]/2,0];
            	            this.inputs[0].dir = LiteGraph.UP;
            	            break;
            	        case "DOWN":
							this.inputs[0].pos = [this.size[0]/2,this.size[1]];
            	            this.inputs[0].dir = LiteGraph.DOWN;
            	            break;
            	    }
            	    switch (o)
            	    {
            	        case "LEFT":
            	            this.outputs[0].pos = [0,0.7*LiteGraph.NODE_SLOT_HEIGHT+LiteGraph.CANVAS_GRID_SIZE];
            	            this.outputs[0].dir = LiteGraph.LEFT;
            	            break;
            	        case "RIGHT":
            	        	this.outputs[0].pos = [this.size[0],0.7*LiteGraph.NODE_SLOT_HEIGHT+LiteGraph.CANVAS_GRID_SIZE];
            	            this.outputs[0].dir = LiteGraph.RIGHT;
            	            break;
            	        case "UP":
            	        	this.outputs[0].pos = [this.size[0]/2,0];
            	            this.outputs[0].dir = LiteGraph.UP;
            	            break;
            	        case "DOWN":
							this.outputs[0].pos = [this.size[0]/2,this.size[1]];
            	            this.outputs[0].dir = LiteGraph.DOWN;
            	            break;
            	    }
				}

				this.onConnectionsChange = function (type, index, connected, link_info) {
					if (this.configured && connected && type === LiteGraph.OUTPUT) {
						const types = new Set(this.outputs[0].links.map((l) => app.graph.links[l].type).filter((t) => t !== "*"));
						if (types.size > 1) {
							const linksToDisconnect = [];
							for (let i = 0; i < this.outputs[0].links.length - 1; i++) {
								const linkId = this.outputs[0].links[i];
								const link = app.graph.links[linkId];
								linksToDisconnect.push(link);
							}
							for (const link of linksToDisconnect) {
								const node = app.graph.getNodeById(link.target_id);
								node.disconnectInput(link.target_slot);
							}
						}
					}

					let currentNode = this;
					let updateNodes = [];
					let inputType = null;
					let inputNode = null;
					while (currentNode) {
						updateNodes.unshift(currentNode);
						const linkId = currentNode.inputs[0].link;
						if (linkId !== null) {
							const link = app.graph.links[linkId];
							if (!link) return;
							const node = app.graph.getNodeById(link.origin_id);
							const type = node.constructor.type;
							if (type.includes ("Reroute")) {
								if (node === this) {
									currentNode.disconnectInput(link.target_slot);
									currentNode = null;
								} else {
									currentNode = node;
								}
							} else {
								inputNode = currentNode;
								inputType = node.outputs[link.origin_slot]?.type ?? null;
								break;
							}
						} else {
							currentNode = null;
							break;
						}
					}

					const nodes = [this];
					let outputType = null;
					while (nodes.length) {
						currentNode = nodes.pop();
						const outputs = (currentNode.outputs ? currentNode.outputs[0].links : []) || [];
						if (outputs.length) {
							for (const linkId of outputs) {
								const link = app.graph.links[linkId];

								if (!link) continue;

								const node = app.graph.getNodeById(link.target_id);
								const type = node.constructor.type;

								if (type.includes("Reroute")) {
									nodes.push(node);
									updateNodes.push(node);
								} else {
									const nodeOutType =
										node.inputs && node.inputs[link?.target_slot] && node.inputs[link.target_slot].type
											? node.inputs[link.target_slot].type
											: null;
									if (this.configured && inputType && inputType !== "*" && nodeOutType !== inputType) {
										node.disconnectInput(link.target_slot);
									} else {
										outputType = nodeOutType;
									}
								}
							}
						}
					}

					this.linkType = inputType || outputType || "*";
					const linkColor = LGraphCanvas.link_type_colors[this.linkType];

					let widgetConfig;
					let targetWidget;
					let widgetType;
					for (const node of updateNodes) {
						node.outputs[0].type = inputType || "*";
						node.__outputType = this.linkType;
						node.outputs[0].name = node.properties.showOutputText ? this.linkType : "";
						if (node.linkType) node.linkType = this.linkType;
						node.size = node.computeSize();
						if (node.applyOrientation) node.applyOrientation();

						for (const l of node.outputs[0].links || []) {
							const link = app.graph.links[l];
							if (link) {
								link.color = linkColor;

								if (app.configuringGraph) continue;
								const targetNode = app.graph.getNodeById(link.target_id);
								const targetInput = targetNode.inputs?.[link.target_slot];
								if (targetInput?.widget) {
									const config = getWidgetConfig(targetInput);
									if (!widgetConfig) {
										widgetConfig = config[1] ?? {};
										widgetType = config[0];
									}
									if (!targetWidget) {
										targetWidget = targetNode.widgets?.find((w) => w.name === targetInput.widget.name);
									}

									const merged = mergeIfValid(targetInput, [config[0], widgetConfig]);
									if (merged.customConfig) {
										widgetConfig = merged.customConfig;
									}
								}
							}
						}
					}

					for (const node of updateNodes) {
						if (widgetConfig && outputType) {
							node.inputs[0].widget = { name: "value" };
							setWidgetConfig(node.inputs[0], [widgetType ?? this.linkType, widgetConfig], targetWidget);
						} else {
							setWidgetConfig(node.inputs[0], null);
						}
					}

					if (inputNode) {
						const link = app.graph.links[inputNode.inputs[0].link];
						if (link) {
							link.color = linkColor;
						}
					}
				};

				this.clone = function () {
					const cloned = MxRerouteNode.prototype.clone.apply(this);
					cloned.removeOutput(0);
					cloned.addOutput("", "*");
					cloned.size = cloned.computeSize();
					cloned.ioOrientation();
					return cloned;
				};

				this.isVirtualNode = true;
			}

			computeSize() { return [2*LiteGraph.CANVAS_GRID_SIZE+1.4*LiteGraph.NODE_SLOT_HEIGHT,2*LiteGraph.CANVAS_GRID_SIZE+1.4*LiteGraph.NODE_SLOT_HEIGHT] }
		}

		LiteGraph.registerNodeType(
			"mxReroute",
			Object.assign(MxRerouteNode, {
				title_mode: LiteGraph.NO_TITLE,
				title: "mxReroute",
				collapsable: false,
			})
		);

		MxRerouteNode.category = "utils";
	},
});
